#include "stdio.h"
#include "string.h"
#include "rbtree.h"


struct mytype {
      struct rb_node node;
      char *keystring;
};

struct mytype *my_search(struct rb_root *root, char *string)
{
      struct rb_node *node = root->rb_node;

      while (node) {
              struct mytype *data = container_of(node, struct mytype, node);
              int result;

              result = strcmp(string, data->keystring);

              if (result < 0)
                      node = node->rb_left;
              else if (result > 0)
                      node = node->rb_right;
              else
                      return data;
      }
      return NULL;
}


int my_insert(struct rb_root *root, struct mytype *data)
{
      struct rb_node **new = &(root->rb_node), *parent = NULL;

      /* Figure out where to put new node */
      while (*new) {
              struct mytype *this = container_of(*new, struct mytype, node);
              int result = strcmp(data->keystring, this->keystring);

              parent = *new;
              if (result < 0)
                      new = &((*new)->rb_left);
              else if (result > 0)
                      new = &((*new)->rb_right);
              else
                      return FALSE;
      }

      /* Add new node and rebalance tree. */
      rb_link_node(&data->node, parent, new);
      rb_insert_color(&data->node, root);

      return TRUE;
}


int main()
{
    struct rb_root root = RB_ROOT;
    // struct mytype* data1 = (struct mytype*)malloc(sizeof(struct mytype));
    // struct mytype* data2 = (struct mytype*)malloc(sizeof(struct mytype));
    // struct mytype* data3 = (struct mytype*)malloc(sizeof(struct mytype));
    struct mytype data1, data2, data3;
    rb_init_node(&data1.node);
    rb_init_node(&data2.node);
    rb_init_node(&data3.node);
    data1.keystring = "hello";
    data2.keystring = "world";
    data3.keystring = "baichen";

    my_insert(&root, &data1);
    my_insert(&root, &data2);
    my_insert(&root, &data3);

    struct rb_node* cur = rb_first(&root);
    int i = 0;
    printf("before delete\n");
    while (1) {
        struct mytype* data = container_of(cur, struct mytype, node);
        printf("data%d: %s\n", i, data->keystring);
        i++;
        if (!(cur = rb_next(cur))) break;
    }
    struct mytype* ret = my_search(&root, "baichen");
    printf("data=%s\nstart delete\n", ret->keystring);
    rb_erase(&ret->node, &root);
    printf("delete finished\n");
    i = 0;

    cur = rb_first(&root);
    while (1) {
        struct mytype* data = container_of(cur, struct mytype, node);
        printf("data%d: %s\n", i, data->keystring);
        i++;
        if (!(cur = rb_next(cur))) break;
    }
    return 0;
}